import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import make_blobs
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
import time
from model import SVM


class MultiClassSVM:
    def __init__(self, max_iter=100, kernel_type='linear', C=1.0, epsilon=0.001):
        """
        初始化多分类SVM模型（一对多策略）
        :param max_iter: 最大迭代次数
        :param kernel_type: 核函数类型，可选 'linear', 'poly', 'rbf'
        :param C: 惩罚参数
        :param epsilon: 精度
        """
        self.max_iter = max_iter
        self.kernel_type = kernel_type
        self.C = C
        self.epsilon = epsilon
        self.models = []
        self.classes = None

    def fit(self, X, y):
        """
        训练多分类SVM模型
        :param X: 训练数据特征矩阵
        :param y: 训练数据标签向量
        """
        self.classes = np.unique(y)
        print(f"训练多分类SVM，类别数: {len(self.classes)}")

        # 为每个类别训练一个SVM模型（一对多策略）
        for c in self.classes:
            print(f"训练类别 {c} 的SVM模型...")
            # 创建二分类标签: 当前类为1，其他类为-1
            y_binary = np.where(y == c, 1, -1)

            # 创建并训练SVM
            svm = SVM(max_iter=self.max_iter, kernel_type=self.kernel_type,
                      C=self.C, epsilon=self.epsilon)
            svm.fit(X, y_binary)
            self.models.append(svm)

        return self

    def predict(self, X):
        """
        预测新数据的类别
        :param X: 特征矩阵
        :return: 预测类别
        """
        n_samples = X.shape[0]
        # 存储每个模型的决策函数值
        decision_values = np.zeros((n_samples, len(self.classes)))

        # 计算每个样本在每个模型中的决策函数值
        for i, model in enumerate(self.models):
            # 使用SVM的predict获取原始决策值（不使用sign函数）
            if model.kernel_type == 'linear' and model.w is not None:
                decision_values[:, i] = np.dot(X, model.w) + model.b
            else:
                for j in range(n_samples):
                    s = 0
                    for alpha, sv_y, sv in zip(model.alpha, model.support_vector_labels, model.support_vectors):
                        s += alpha * sv_y * \
                            model.kernels[model.kernel_type](X[j], sv)
                    decision_values[j, i] = s + model.b

        # 返回决策函数值最大的类别的索引（对应于self.classes中的类别）
        return self.classes[np.argmax(decision_values, axis=1)]

    def score(self, X, y):
        """计算准确率"""
        y_pred = self.predict(X)
        return np.mean(y_pred == y)


def generate_data(n_samples=300, n_classes=3, n_features=2, random_state=42):
    """
    生成多分类数据
    :param n_samples: 样本数量
    :param n_classes: 类别数量
    :param n_features: 特征数量
    :param random_state: 随机种子
    :return: X - 特征矩阵, y - 标签
    """
    X, y = make_blobs(n_samples=n_samples, centers=n_classes,
                      n_features=n_features, random_state=random_state)

    # 标准化特征
    scaler = StandardScaler()
    X = scaler.fit_transform(X)

    return X, y


def plot_data_and_decision_boundary(X, y, model, title='Multi-Class SVM Decision Boundary'):
    """
    绘制数据点和决策边界
    :param X: 特征矩阵（二维）
    :param y: 标签
    :param model: 训练好的MultiClassSVM模型
    :param title: 图表标题
    """
    # 仅对二维数据有效
    if X.shape[1] != 2:
        print("只能绘制二维数据的决策边界")
        return

    plt.figure(figsize=(12, 10))

    # 创建网格以绘制决策边界
    x_min, x_max = X[:, 0].min() - 1, X[:, 0].max() + 1
    y_min, y_max = X[:, 1].min() - 1, X[:, 1].max() + 1
    xx, yy = np.meshgrid(np.arange(x_min, x_max, 0.02),
                         np.arange(y_min, y_max, 0.02))

    # 进行预测
    Z = model.predict(np.c_[xx.ravel(), yy.ravel()])
    Z = Z.reshape(xx.shape)

    # 绘制决策边界
    plt.contourf(xx, yy, Z, alpha=0.3, cmap=plt.cm.viridis)
    plt.contour(xx, yy, Z, colors='k', linestyles='-', linewidths=1)

    # 绘制数据点
    unique_classes = np.unique(y)
    colors = plt.cm.viridis(np.linspace(0, 1, len(unique_classes)))

    for i, cls in enumerate(unique_classes):
        plt.scatter(X[y == cls, 0], X[y == cls, 1],
                    color=colors[i], s=40, edgecolor='k',
                    label=f'Class {cls}')

    # 绘制支持向量
    for i, model in enumerate(model.models):
        plt.scatter(model.support_vectors[:, 0], model.support_vectors[:, 1],
                    s=100, facecolors='none', edgecolors='r', alpha=0.5)

    plt.title(title)
    plt.xlabel('Feature 1')
    plt.ylabel('Feature 2')
    plt.legend()
    plt.tight_layout()
    plt.show()


def test_multi_class_svm_with_different_kernels():
    """测试多分类SVM在不同核函数下的性能"""
    # 生成数据
    X, y = generate_data(n_samples=300, n_classes=3, n_features=2)

    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42)

    # 不同的核函数
    kernels = ['linear', 'poly', 'rbf']

    for kernel in kernels:
        print(f"\n使用 {kernel} 核函数测试多分类SVM")

        # 创建并训练模型
        start_time = time.time()
        model = MultiClassSVM(
            max_iter=100, kernel_type=kernel, C=1.0, epsilon=0.001)
        model.fit(X_train, y_train)
        training_time = time.time() - start_time

        # 计算准确率
        train_acc = model.score(X_train, y_train)
        test_acc = model.score(X_test, y_test)

        print(f"训练时间: {training_time:.2f}秒")
        print(f"训练集准确率: {train_acc:.4f}")
        print(f"测试集准确率: {test_acc:.4f}")

        # 可视化决策边界
        plot_data_and_decision_boundary(X, y, model,
                                        title=f'Multi-Class SVM with {kernel.upper()} Kernel')


if __name__ == "__main__":
    # 测试多分类SVM
    test_multi_class_svm_with_different_kernels()

    # 更多类别的测试
    X, y = generate_data(n_samples=500, n_classes=5,
                         n_features=2, random_state=42)
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42)

    print("\n测试5个类别的分类")
    model = MultiClassSVM(max_iter=100, kernel_type='rbf',
                          C=1.0, epsilon=0.001)
    model.fit(X_train, y_train)

    train_acc = model.score(X_train, y_train)
    test_acc = model.score(X_test, y_test)

    print(f"训练集准确率: {train_acc:.4f}")
    print(f"测试集准确率: {test_acc:.4f}")

    plot_data_and_decision_boundary(
        X, y, model, title='Multi-Class SVM (5 Classes)')
